#!/usr/bin/env perl

# Copyright (C) Yichun Zhang (agentzh)

use 5.006001;
use strict;
use warnings;

use Getopt::Std qw( getopts );

my %opts;

getopts("a:dhp:n:", \%opts)
    or die usage();

if ($opts{h}) {
    print usage();
    exit;
}

my $pid = $opts{p}
    or die "No nginx process pid specified by the -p option\n";

if ($pid !~ /^\d+$/) {
    die "Bad -p option value \"$pid\": not look like a pid\n";
}

my $stap_args = $opts{a} || '';

if ($^O ne 'linux') {
    die "Only linux is supported but I am on $^O.\n";
}

my $exec_file = "/proc/$pid/exe";
if (!-f $exec_file) {
    die "Nginx process $pid is not running or ",
        "you do not have enough permissions.\n";
}

my $nginx_path = readlink $exec_file;

my $ver = `stap --version 2>&1`;
if (!defined $ver) {
    die "Systemtap not installed or its \"stap\" utility is not visible to the PATH environment: $!\n";
}

if ($ver =~ /version\s+(\d+\.\d+)/i) {
    my $v = $1;
    if ($v < 2.1) {
        die "ERROR: at least systemtap 2.1 is required but found $v\n";
    }

} else {
    die "ERROR: unknown version of systemtap:\n$ver\n";
}

my $stap_src;

my $zone = $opts{n};

my $preamble = <<_EOC_;
probe begin {
    printf("Tracing %d ($nginx_path)...\\n\\n", target())
}
_EOC_
chop $preamble;

if ($zone) {
    my $quoted_zone = quote_str($zone);

    my $zone_name_len = length $zone;

    $stap_src = <<_EOC_;
$preamble

probe process("$nginx_path").function("ngx_process_events_and_timers"),
    process("$nginx_path").function("ngx_http_init_request")!,
    process("$nginx_path").function("ngx_http_create_request")
{
    if (pid() == target()) {

        found = 0
        begin = local_clock_us()

        pagesize = \@var("ngx_pagesize\@ngx_alloc.c")
        part = &\@var("ngx_cycle\@ngx_cycle.c")->shared_memory->part
        zone = \@cast(part, "ngx_list_part_t")->elts
        for (i = 0; ; i++) {
            if (i >= \@cast(part, "ngx_list_part_t")->nelts) {
                if (\@cast(part, "ngx_list_part_t")->next == 0) {
                    break
                }

                part = \@cast(part, "ngx_list_part_t")->next
                zone = \@cast(part, "ngx_list_part_t")->elts
                i = 0
            }

            shm = &\@cast(zone, "ngx_shm_zone_t")[i]->shm

            name = &\@cast(shm, "ngx_shm_t")->name

            if (\@cast(name, "ngx_str_t")->len != $zone_name_len) {
                continue;
            }

            zone_name = user_string_n(\@cast(name, "ngx_str_t")->data, $zone_name_len)

            if (zone_name != $quoted_zone) {
                continue;
            }

            init = \@cast(zone, "ngx_shm_zone_t")[i]->init
            addr = \@cast(shm, "ngx_shm_t")->addr
            addr2 = \@cast(addr, "ngx_slab_pool_t")->addr

            if (addr != addr2) {
                error("shm zone \\"" . zone_name . "\\" is corrupted or "
                      . "not yet loaded into the process memory space: "
                      . "pool != pool->addr "
                      . sprintf("(pool->addr: %p)", addr2))
                exit()
            }

            init_name = usymname(init)
            owner = str_replace(str_replace(init_name, "_init", ""), "_zone", "")

            printf("shm zone \\"%s\\"\\n", zone_name)
            printf("    owner: %s\\n", owner)
            printf("    total size: %d KB\\n", \@cast(shm, "ngx_shm_t")->size / 1024)

            pages = 0
            free = &\@cast(addr, "ngx_slab_pool_t")->free
            blocks = 0

            /*
            printf("shm: %p, shm->addr: %p, free: %p, addr2: %p, next: %p\\n",
                   shm, addr, free, addr2,
                   \@cast(free, "ngx_slab_page_t")->next)
            */

            for (page = \@cast(free, "ngx_slab_page_t")->next;
                 page && page != free;
                 page = \@cast(page, "ngx_slab_page_t")->next)
            {
                //printf("page: %p\\n", page)
                pages += \@cast(page, "ngx_slab_page_t")->slab
                if (++blocks % 1000 == 0) {
                    printf("\\r    free pages: %d KB (%d pages, %d blocks)",
                           pages * pagesize / 1024, pages, blocks)
                }
            }

            printf("\\r    free pages: %d KB (%d pages, %d blocks)\\n",
                   pages * pagesize / 1024, pages, blocks)

            found = 1
            break
        }

        if (!found) {
            printf("shm zone \\"%s\\" not found.\\n", $quoted_zone)
        }

        elapsed = local_clock_us() - begin
        printf("\\n%d microseconds elapsed in the probe handler.\\n", elapsed)

        exit()

    } /* pid() == target() */
}
_EOC_

} else {
    # no zone name specified

    $stap_src = <<_EOC_;
$preamble

probe process("$nginx_path").function("ngx_process_events_and_timers"),
    process("$nginx_path").function("ngx_http_handler")
{
    if (pid() == target()) {

        begin = local_clock_us()

        part = &\@var("ngx_cycle\@ngx_cycle.c")->shared_memory->part
        zone = \@cast(part, "ngx_list_part_t")->elts
        for (i = 0; ; i++) {
            if (i >= \@cast(part, "ngx_list_part_t")->nelts) {
                if (\@cast(part, "ngx_list_part_t")->next == 0) {
                    break
                }

                part = \@cast(part, "ngx_list_part_t")->next
                zone = \@cast(part, "ngx_list_part_t")->elts
                i = 0
            }

            shm = &\@cast(zone, "ngx_shm_zone_t")[i]->shm
            addr = \@cast(shm, "ngx_shm_t")->addr
            init = \@cast(zone, "ngx_shm_zone_t")[i]->init
            name = &\@cast(shm, "ngx_shm_t")->name

            zone_name = user_string_n(\@cast(name, "ngx_str_t")->data,
                                      \@cast(name, "ngx_str_t")->len)

            addr2 = \@cast(addr, "ngx_slab_pool_t")->addr

            if (addr != addr2) {
                println("WARNING: shm zone \\"" . zone_name . "\\" is "
                      . "corrupted or "
                      . "not yet loaded into the process memory space: "
                      . "pool != pool->addr "
                      . sprintf("(pool->addr: %p)", addr2))
                exit()
            }

            init_name = usymname(init)
            owner = str_replace(str_replace(init_name, "_init", ""), "_zone", "")

            printf("shm zone \\"%s\\"\\n    owner: %s\\n    total size: %d KB\\n\\n",
                   user_string_n(\@cast(name, "ngx_str_t")->data,
                                 \@cast(name, "ngx_str_t")->len),
                   owner, \@cast(shm, "ngx_shm_t")->size / 1024)
        }

        println("Use the -n <zone> option to see more details about each zone.")

        elapsed = local_clock_us() - begin
        printf("%d microseconds elapsed in the probe.\\n", elapsed)
        exit()
    } /* pid() == target() */
}
_EOC_
}

if ($opts{d}) {
    print $stap_src;
    exit;
}

if (!$zone) {
    $stap_args .= " --skip-badvars"
}

open my $in, "|stap $stap_args -x $pid -"
    or die "Cannot run stap: $!\n";

print $in $stap_src;

close $in;

sub usage {
    return <<'_EOC_';
Usage:
    ngx-shm [optoins]

Options:
    -a <args>           Pass extra arguments to the stap utility.
    -d                  Dump out the systemtap script source.
    -h                  Print this usage.
    -n <zone>           Show details about the shm zone named <zone>.
    -p <pid>            Specify the nginx worker process pid.

Examples:
    ngx-shm -p 12345
    ngx-shm -p 12345 -n dogs
    ngx-shm -p 12345 -n my_zone_name -a '-DMAXACTION=100000'
_EOC_
}

sub quote_str {
    my $s = shift;
    $s =~ s/\\/\\\\/g;
    $s =~ s/"/\\"/g;
    $s =~ s/\n/\\n/g;
    $s =~ s/\t/\\t/g;
    $s =~ s/\r/\\r/g;
    return qq{"$s"};
}

